#include "inference.h"
#include <iostream>
#include <fstream>
#include <opencv2/imgcodecs.hpp>
#include <nlohmann/json.hpp>

using json = nlohmann::json;

int main(int argc, char **argv) {
	// Check if the correct number of arguments is provided
	if (argc != 4) {
		std::cerr << "usage: " << argv[0] << " <model_path> <image_path> <output_json_path>" << std::endl;
		return 1;
	}

	// Get the model, image paths, and output path from the command-line arguments
	const std::string model_path = argv[1];
	const std::string image_path = argv[2];
	const std::string output_path = argv[3];

	// Read the input image
	cv::Mat image = cv::imread(image_path);

	// Check if the image was successfully loaded
	if (image.empty()) {
		json error_response;
		error_response["error"] = "Failed to load image";
		error_response["image_path"] = image_path;
		
		std::ofstream output_file(output_path);
		output_file << error_response.dump(4);
		output_file.close();
		return 1;
	}

	// Define the confidence and NMS thresholds
	const float confidence_threshold = 0.5;
	const float NMS_threshold = 0.5;

	try {
		// Initialize the YOLO inference with the specified model and parameters
		yolo::Inference inference(model_path, cv::Size(640, 640), confidence_threshold, NMS_threshold);

		// Run inference on the input image
		inference.RunInference(image);

		// Get model configuration
		yolo::ModelConfig config = inference.GetModelConfig();

		// Convert model version to string
		std::string version_str;
		switch (config.version) {
			case yolo::ModelVersion::YOLOv8:
				version_str = "YOLOv8";
				break;
			case yolo::ModelVersion::YOLOv11:
				version_str = "YOLOv11";
				break;
			default:
				version_str = "Unknown";
				break;
		}

		// Get device information
		std::string device_info = inference.GetDeviceInfo();

		// Prepare JSON response with detection results
		json response;
		response["status"] = "success";
		response["image_path"] = image_path;
		response["model_path"] = model_path;
		response["model_version"] = version_str;
		response["num_classes"] = config.num_classes;
		response["image_width"] = image.cols;
		response["image_height"] = image.rows;
		response["device"] = device_info;
		response["detections"] = json::array();

		// Write response to file
		std::ofstream output_file(output_path);
		output_file << response.dump(4);
		output_file.close();

		return 0;
	} catch (const std::exception &e) {
		json error_response;
		error_response["error"] = e.what();
		error_response["status"] = "failed";
		
		std::ofstream output_file(output_path);
		output_file << error_response.dump(4);
		output_file.close();
		return 1;
	}
}

